from textwrap import indent
import requests, json, re

res = requests.post('https://idf-api-prod.aibs-idk-prod.net/', json = {
    "operationName": "Taxonomy",
    "variables": {
        "dataset": "aibs_mouse_ctx-hpf_10x"
    },
    "query": '''
        query Taxonomy($dataset: String) {
            getTranscriptomicDataSet(DataSet: $dataset) {
                cellTypeTaxonomy {
                    taxonomy: nodes {
                        id: accessionId parentId color order alias label __typename
                    }
                    __typename
                }
                __typename
            }
        }
    '''
})

cellTypes = [cellType['alias'] for cellType in json.loads(res.text)['data']['getTranscriptomicDataSet']['cellTypeTaxonomy'][0]['taxonomy'] if cellType['alias'] and re.match(r'^[0-9]+_', cellType['alias'])]
# print(cellTypes)

genes = [ "Pax6", "Neurog2"]

res = requests.post('https://idf-api-prod.aibs-idk-prod.net/', json = {
    "variables": {
        "rows": cellTypes,
        "features": genes,
        "dataset": "aibs_mouse_ctx-hpf_10x",
        "operator": "TRIMMED_MEANS",
        "groupBy": "CLUSTER_LABEL"
    },
    "query": '''
        query (
            $dataset: String,
            $features: [String],
            $rows: [String],
            $operator: MatrixAggregationOperator,
            $groupBy: MatrixAggregationCellMetadata
        )
        {
            aggregateRowsOnFeatureMatrix(
                dataset: $dataset,
                features: $features,
                rows: $rows,
                operator: $operator,
                groupBy: $groupBy
            )
            {
                groupByResults {
                    row
                    featureResults {
                        feature
                        value
                        __typename
                    }
                    __typename
                }
                __typename
            }
        }
    '''
})

# print(res.text)

expressions = json.loads(res.text)['data']['aggregateRowsOnFeatureMatrix']['groupByResults']

with open('expressions.json', 'w', encoding='UTF-8') as file:
    file.write(json.dumps(expressions, indent = 4))

import xlwt
f = xlwt.Workbook() #创建工作簿
sheet = f.add_sheet(u'expressions',cell_overwrite_ok=True) #创建sheet

# 生成第一行
rows = ['cell']
rows.extend(genes)

maps = dict(zip(rows, range(len(rows))))

for i in range(len(rows)):
    sheet.write(0, i, rows[i])

for i in range(len(expressions)):
    sheet.write(i+1, 0, expressions[i]['row'])
    for j in range(len(expressions[i]['featureResults'])):
        sheet.write(
            i+1,
            maps[expressions[i]['featureResults'][j]['feature']],
            expressions[i]['featureResults'][j]['value']
        )

f.save('expressions.xls')
